import copy

class BaseGate:
    def __init__(self, duration, latency=0, awg_flag = None, pulse_type = None, amp = None, freq = None, phase = None, label = None):
        self.awg_flag = False
        self.on_flag = False
        self.label = None
        self.ion_index = None
        self.duration = duration
        self.latency = latency
        self.pulse_type = pulse_type
        if awg_flag != None:
            self.awg_flag = awg_flag
        if label != None:
            self.label = label
        self.amp = amp
        self.freq = freq
        self.phase = phase
        self.on_flag = False
        self.segment_flag = False
        self.segment_number = -1
        self.composite_flag = False

    def on(self, ion_index):
        assert isinstance(ion_index, int) or isinstance(ion_index, list)
        #self.ion_index = ion_index
        if self.on_flag == False:
            self.ion_index = ion_index
            self.on_flag = True
            return self
        else:
            another_gate = copy.copy(self)
            another_gate.ion_index = ion_index
            return another_gate

    def set_pulse_parameters(self, duration = None, awg_flag = None, amp = None, freq = None, phase = None):
        if duration != None:
            self.duration = duration
        if awg_flag != None:
            self.awg_flag = awg_flag
        if amp != None:
            self.amp = amp
        if freq != None:
            self.freq = freq
        if phase != None:
            self.phase = phase

    def parameter_list(self, segment_id = None):
        if segment_id == None:
            return {'amp':self.amp, 'freq':self.freq, 'phase':self.phase, 'duration':self.duration}
        else:
            #print({'amp':self.amp[segment_id], 'freq':self.freq[segment_id], 'phase':self.phase[segment_id], 'duration':self.duration})
            return {'amp':self.amp[segment_id], 'freq':self.freq[segment_id], 'phase':self.phase[segment_id], 'duration':self.duration}
        # self.para_set = True

    def update_hardware(self):
        '''
        you should clarify the concrete hardware interface for each subclass
        See ZTESTER_Load_ion.py for an example
        '''
        raise AssertionError("'"+self.label+".update_hardware' is not DEFINED.")

    def reset_hardware(self):
        raise AssertionError("'"+self.label+".reset_hardware' is not DEFINED.")
        

class AWGBaseGate(BaseGate):
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None,\
                 label = None, segment_number = -1, time_intervals = None):
        super().__init__(duration, latency, awg_flag, 'Rx', amp, freq, phase, label)
        self.segment_flag = False
        self.time_intervals = None
        if segment_number > 0:
            self.segment_flag = True
            self.segment_number = segment_number
            self.time_intervals = time_intervals
            assert isinstance(self.time_intervals, list)
            for pair in self.time_intervals:
                assert len(pair) == 2
            assert len(self.time_intervals) == self.segment_number


class AdvanceGate:
    def __init__(self, duration, latency=0, ion_index = None, para_table = {}, pulse_type = None, awg_flag = True):
        self.duration = duration
        self.latency = latency
        self.ion_index = ion_index
        self.para_table = para_table
        self.pulse_type = pulse_type
        self.awg_flag = awg_flag
        self.ion_number = 0
        if self.para_table != {}:
            self.__check_dict(self.para_table)
            self.ion_number = self.para_table['ion_number']
            #print(self.ion_number)

    def __check_dict(self, para_dic):
        '''
        para_dic is a dict with following:
        'ion_number' : int
        'segment_number' : int
        'time_intervals' : list of tuples, like -- [(t1,t2),(t3,t4),(t5,t6),...]
        'data_per_ion' : dict of dicts, each dict contains three keys 'amp', 'freq', 'phase'
        '''
        assert type(para_dic.get('ion_number')) == int
        assert type(para_dic.get('segment_number')) == int
        ion_number = para_dic['ion_number']
        seg_number = para_dic['segment_number']
        assert ion_number == len(para_dic['data_per_ion'])
        for (ion,ion_data) in para_dic['data_per_ion'].items():
            #print(ion,ion_data)
            if seg_number > 1:
                assert len(ion_data['amp']) == seg_number
                assert len(ion_data['freq']) == seg_number
                assert len(ion_data['phase']) == seg_number

    def set_parameter_table(self, para_dic):
        self.__check_dict(para_dic)
        self.para_table = para_dic
        self.ion_number = para_dic['ion_number']

    def on(self, ion_index):
        assert type(ion_index) in {int, list}

        if type(ion_index) == list:
            if len(ion_index) > self.ion_number:
                print('{} has more qubits than ion_number({}) in this advanced gate!'.format(ion_index, self.ion_number))
                raise SyntaxError()

        def iter_int_or_tuple(item):
            try:
                for sub_item in item:
                    yield sub_item
            except TypeError:
                yield item

        for ion in iter_int_or_tuple(ion_index):
            if self.para_table['data_per_ion'].get(ion) == None:
                print('ion index {} not existing in this advanced gate!'.format(ion))
            #self.para_table['data_per_ion'][ion] = 0
        return ('advanced gate', ion_index, self)

class parallel_entanglement(AdvanceGate):
        def __init__(self, duration, latency=0, para_table = {}, pulse_type = 'xx', awg_flag = True):
            super().__init__(duration, latency=latency, para_table=para_table)
            self.pulse_type = 'parallel_' + pulse_type


def sync(ion_index):
    # synchronize all ions in ion_index, aligning their timing
    if type(ion_index) == int:
        ion_index = [ion_index]
    return ('sync', ion_index)

#awg_trigger函数用于打开awg开关，latency参数表示距离上一个输入波形结束时的延时
def awg_trigger(latency = 0, ion_index = None):
    # turn on the awg trigger, latency is the time distance with previous pulse
    if type(ion_index) == int:
        ion_index = [ion_index]
    return ('awg_trigger', latency, ion_index)

#def pending_all(latency = 0):
def pending(latency = 0):
    # synchoronize the timestamp of all ions, which means the following pulses start at the time of the lastest pulse
    # for example, if the maxmimal ending time of previous pulses is 'tm', then after a pending \
        #  operation the following pulses will start at 'tm+lantency'
    # default pending latency is 0
    return ('pending', latency)
